from typing import Optional

from pydantic import BaseModel

from app.models.plans import StudyPlanStatus
from app.schemas.studies import MedicalStudySchema


class StudyPlanCreateSchema(BaseModel):
    year: int
    month: int
    study_id: int
    number_of_studies: int
    status: Optional[StudyPlanStatus] = StudyPlanStatus.FORECASTED


class StudyPlanUpdateSchema(BaseModel):
    id: int
    number_of_studies: int
    status: Optional[StudyPlanStatus] = StudyPlanStatus.FORECASTED


class StudyPlanSchema(BaseModel):
    id: int
    year: int
    month: int
    study: MedicalStudySchema
    number_of_studies: int
    status: StudyPlanStatus
